import pandas as pd
import numpy as np
import tushare as ts
from pymongo import MongoClient
import yaml


def loc_db(db_name, cfg={}):
    if len(cfg) == 0:
        cfg = load_config_()
    return MongoClient(cfg['mongodb'])[db_name]


def load_config_():
    with open("../config.yaml", 'r', encoding='utf-8') as f:
        data = yaml.load(f.read(), yaml.FullLoader)
    print("读取本地配置")
    return data


def loc_collection(db_name, collection_name, cfg={}):
    if len(cfg) == 0:
        cfg = load_config_()
    conn = MongoClient(cfg['mongodb'])
    coll_ = conn[db_name][collection_name]
    return coll_


def to_week(df):
    df.index = [pd.to_datetime(str(s)) for s in df.index]
    # idx = df.index.to_list()
    df['year'] = df.copy().index.isocalendar().year
    df['week'] = df.copy().index.isocalendar().week
    df['wk'] = [f"{a}_{b}" for a, b in zip(df['year'], df['week'])]
    df.index = df['wk']
    return df.iloc[:, 0]


def drop_dup_idx(li):
    new_li = []
    for df in li:
        idx = df.index.to_list()
        if len(idx) == len(set(idx)):
            new_li.append(df)
    return new_li


def week_to_date(t):
    year, week = t.split('_')
    t = pd.to_datetime(f"{year}-12-31")
    w = t.isocalendar()[1]
    while str(w) != str(week):
        t = t - pd.Timedelta(days=1)
        w = t.isocalendar()[1]
    return t


def get_weekly_daily_data(freq='周'):
    cfg = load_config_()
    db = loc_db('to_it', cfg)
    coll_names = db.list_collection_names()
    coll_names = [s for s in coll_names if len(s) > 7]
    c_name = sorted(coll_names)[-1]
    coll = loc_collection('to_it', c_name)
    data = list(coll.find({'freq': '周'}))
    filter_data = []
    for dic in data:
        columns = dic['columns']
        if dic['mode'] in ['期货', '基差', '价差'] and '主力' in columns[0]:
            filter_data.append(dic)
        elif dic['mode'] in ['现货', '库存'] and dic['default'] == True:
            filter_data.append(dic)
    df_li = [pd.DataFrame(d['value'], index=d['t'], columns=d['columns']) for d in filter_data]
    se_li = [df.iloc[:, [0]] for df in df_li]
    len_se_li = [s for s in se_li if len(s) > np.mean([len(s) for s in se_li])]
    if freq == '周':
        idx_se_li = [to_week(s)[-300:] for s in len_se_li]
    elif freq == '日':
        idx_se_li = [s[-500:] for s in len_se_li]
    idx_se_li = drop_dup_idx(idx_se_li)
    data = pd.concat(idx_se_li, axis=1)
    if freq == '周':
        data.index = [week_to_date(t) for t in data.index]
    data = data.sort_index()
    cols = data.columns.to_list()
    idx = data.index.to_list()
    if freq == '周':
        for i in range(30):
            idx.append(idx[-1] + pd.Timedelta(days=7))
    else:
        pro = ts.pro_api(cfg['tushare'])
        trade_date = pro.trade_cal(exchange='DCE', start_date='20200101', end_date='20221231')
        filter_trade_date = trade_date[trade_date.is_open == True]
        cal_date = filter_trade_date['cal_date'].to_list()
        ttt = [str(pd.to_datetime(str(s)))[:10] for s in cal_date]
        t_last_i = ttt.index(idx[-1])
        for i in range(t_last_i + 1, t_last_i + 30):
            idx.append(ttt[i])
    data.index = [str(s)[:10] for s in data.index]
    data['date'] = data.index
    new_cols = ['date', *cols]
    data = data.loc[:, new_cols]
    data = data[data.index > '2016-05-01']
    data = data.fillna(method='ffill')
    map_ = {f'a{i}': v for i, v in enumerate(data.columns.to_list())}
    map_df = pd.Series(map_)
    save_cols = ['date']
    for i in range(1, data.shape[1]):
        save_cols.append(f'a{i}')
    data.columns = save_cols
    return data, pd.DataFrame(idx), map_df


def prepare_data(freq='日'):
    data, data_t, map_df = get_weekly_daily_data(freq=freq)
    data.to_excel('future_data.xlsx', index=None)
    print(data.shape, 'data')
    map_df.to_excel('map.xlsx')
    print(map_df.shape, 'map')
    print('t_last', data_t.iloc[:, -1].to_list()[-1])
    if freq == '日':
        data_t.to_excel('trade_date.xlsx')
    else:
        data_t.to_excel('trade_date_w.xlsx')


if __name__ == "__main__":
    """
    为跑informer准备数据，包含期货、现货、基差、价差、库存的数据准备
    """
    prepare_data(freq='周')
